import socket
import time
import struct
import random
from threading import Timer

NIL = 0
ACK = 1
NAK = 2
ACK0 = 3
ACK1 = 4

VER = '3.0'

class DataStream:
    def __init__(self, data: bytes, pktsize: int) -> None:
        self.data = data
        self.length = len(data)
        self.readptr = 0
        self.pktsize = pktsize
        self.amount = self.length // self.pktsize + 1
        pass

    def __iter__(self):
        return self

    def __next__(self):
        if self.readptr < self.length:
            nxtptr = self.readptr + min(self.pktsize, self.length - self.readptr)
            tmpdata = self.data[self.readptr : nxtptr]
            self.readptr = nxtptr
            return tmpdata
        elif self.readptr == self.length:
            self.readptr += 1
            print("void sent")
            return b""
        else:
            print("stopped")
            raise StopIteration

    pass


class RDT30Client:
    def __init__(
        self,
        filename: str,
        pktsize: int,
        target: tuple[str, int],
        errorsize: int,
        timeout: int,
        losssize: int
    ) -> None:
        self.datastream = DataStream(open(filename, "rb").read(), pktsize)
        self.pktsize = pktsize
        self.target = target
        self.clientSocket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        # 0 -> waiting for call 0 from above
        # 1 -> waiting for ack/nak 0
        # 2 -> waiting for call 1 from above
        # 3 -> waiting for ack/nak 1
        self.state = 0
        self.errorsize = errorsize

        self.timeout = timeout
        self.timer = None
        self.losssize = losssize
        pass

    def timercallback(self, sequence: int, data: bytes):
        self.timer.cancel()
        print('timer called at ', self.datastream.readptr)
        self.sendpkt(self.biterror(self.makepkt(data, NIL, sequence)))
        self.timer = Timer(
            interval=self.timeout, function=self.timercallback, args=[sequence, data]
        )
        self.timer.start()
        pass

    def makepkt(self, rawdata: bytes, confirmation: int, sequence: int):
        # rdt 2.1 报文格式:
        # ack/nak 1byte
        # sequence 1byte
        # length 2byte
        # checksum 2byte
        pkt = struct.pack(
            "!BBHH%ds" % len(rawdata), confirmation, sequence, len(rawdata), 0, rawdata
        )
        csum = self.checksum(pkt)
        return struct.pack(
            "!BBHH%ds" % len(rawdata),
            confirmation,
            sequence,
            len(rawdata),
            csum,
            rawdata,
        )

    def sendpkt(self, pkt: bytes):
        if random.randint(1, self.losssize) == 1:
            print('packet loss at ', self.datastream.readptr)
            pass
        else:
            self.clientSocket.sendto(self.biterror(pkt), self.target)

    def checksum(self, data: bytes) -> int:
        csum = 0
        for i in range(1, len(data), 2):
            csum += int.from_bytes(data[i - 1 : i + 1], "big")
        if len(data) % 2 == 1:
            csum += data[-1] << 8 & 0xFF00
        csum = (csum >> 16) + (csum & 0xFFFF)
        csum = (csum >> 16) + (csum & 0xFFFF)
        csum = ~csum & 0xFFFF
        return csum

    def extract(self, pkt: bytes) -> tuple[int, int, int, int, bytes]:
        confirmation, sequence, length, csum = struct.unpack("!BBHH", pkt[:6])
        data = struct.unpack_from("!%ds" % length, pkt, 6)[0]
        return confirmation, sequence, length, csum, data

    def biterror(self, pkt: bytes):
        if random.randint(1, self.errorsize) == 1:
            pos = random.randint(0, len(pkt) - 1)
            off = random.randint(0, 7)
            npkt = bytearray(pkt)
            npkt[pos] ^= 1 << off
            return bytes(npkt)
        return pkt

    def valid(self, pkt: bytes) -> bool:
        return self.checksum(pkt) == 0

    def run(self):
        dataiter = iter(self.datastream)
        partdata = None
        while True:
            try:
                # partdata = next(dataiter)

                match self.state:
                    case 0:
                        partdata = next(dataiter)
                        self.timer = Timer(
                            interval=self.timeout,
                            function=self.timercallback,
                            args=[0, partdata],
                        )
                        print("sq 0 at ", self.datastream.readptr)
                        self.sendpkt(self.makepkt(partdata, NIL, 0))
                        self.timer.start()
                        self.state = 1
                        pass

                    case 1:
                        pktrcvd = self.clientSocket.recvfrom(1024)[0]
                        self.timer.cancel()
                        # self.timer = Timer(
                        #     interval=self.timeout,
                        #     function=self.timercallback,
                        #     args=[0, partdata],
                        # )
                        if self.valid(pktrcvd) and self.extract(pktrcvd)[0] == ACK0:
                            self.state = 2
                        else:
                            print("resend package at ", self.datastream.readptr)
                            self.sendpkt(self.makepkt(partdata, NIL, 0))
                            self.state = 1

                    case 2:
                        partdata = next(dataiter)
                        self.timer = Timer(
                            interval=self.timeout,
                            function=self.timercallback,
                            args=[1, partdata],
                        )
                        print("sq 1 at ", self.datastream.readptr)
                        self.sendpkt(self.makepkt(partdata, NIL, 1))
                        self.timer.start()
                        self.state = 3
                        pass

                    case 3:
                        pktrcvd = self.clientSocket.recvfrom(1024)[0]
                        self.timer.cancel()
                        # self.timer = Timer(
                        #     interval=self.timeout,
                        #     function=self.timercallback,
                        #     args=[1, partdata],
                        # )
                        if self.valid(pktrcvd) and self.extract(pktrcvd)[0] == ACK1:
                            self.state = 0
                        else:
                            print("resend package at ", self.datastream.readptr)
                            self.sendpkt(self.makepkt(partdata, NIL, 1))
                            self.state = 3

                        pass
            except StopIteration:
                break

    pass


c = RDT30Client(
    filename="komeiji.png",
    pktsize=960,
    target=("127.0.0.1", 65432),
    errorsize=100,
    timeout=0.1,
    losssize=100
)
c.run()
